package ru.spbau.database.group3.indexes.btree;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.LinkedList;
import java.util.PriorityQueue;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;

public class BPlusTree<K extends Comparable<K>, V> implements Serializable{

	private int bucketSize;

	private class Node implements Comparable<Node>, Serializable {
		private V value = null;
		private Node parent = null;
		private K key = null;
		private ArrayList<Node> entries; // for Ni contains all Nodes with keys
											// k: Ni.key <= k <= Ni+1.key

		Node() {
			entries = new ArrayList<Node>();
		}

		Node(V value, K key) {
			entries = new ArrayList<Node>();
			this.value = value;
			this.key = key;
		}

		private String padRight(String s, int n) {
			StringBuilder strBld = new StringBuilder();
			for (int i = 0; i < n; ++i) {
				strBld.append(" ");
			}
			return strBld.toString();
		}

		private String constructString(int offset) {
			String res = new String();
			res = padRight(res, offset);
			if (isLeaf()) {
				return res + key + "->" + value + "\n";
			} else {
				res += entries.get(0).key + "\n";
				for (int i = 0; i < entries.size(); ++i) {
					res += entries.get(i).constructString(offset + 2);
				}
			}
			return res;
		}

		public boolean isInternalNode() {
			return entries.size() > 0;
		}

		public boolean isLeaf() {
			return !isInternalNode();
		}

		@Override
		public boolean equals(Object o) {
			if (!(o instanceof BPlusTree.Node)) {
				return false;
			} else {
				return ((Node) o).key.equals(this.key);
			}
		}

		@Override
		public int compareTo(Node o) {
			return (key.compareTo(o.key));
		}

		private Node removeByKey(K key) {
			for (int i = 0; i < entries.size(); ++i) {
				if (entries.get(i).key.equals(key) && entries.get(i).isLeaf()) {
					return entries.remove(i);
				}
			}
			return null;
		}
	}

	// XXX: it's an ugly hack to make things work somehow for now;
	private SortedSet<Node> leafs = new TreeSet<Node>();
	private Node root = null;

	public BPlusTree(int bucketSize) {
		this.bucketSize = bucketSize;
	}

	private Node searchForParentNode(K key) {
		Node currentNode = root;
		Node parent = null;
		while (currentNode != null && !currentNode.isLeaf()) {
			parent = currentNode;
			if (key.compareTo(currentNode.entries.get(0).key) < 0) {
				currentNode = currentNode.entries.get(0);
				continue;
			}
			for (int i = currentNode.entries.size() - 1; i >= 0; --i) {
				if (key.compareTo(currentNode.entries.get(i).key) >= 0) {
					currentNode = currentNode.entries.get(i);
					break;
				}
			}
		}
		return parent;
	}

	private Node searchForNode(K key) {
		Node parent = searchForParentNode(key);
		if (parent == null) {
			if (root != null && root.key.equals(key)) {
				return root;
			} else {
				return null;
			}
		} else {
			for (Node n : parent.entries) {
				if (n.key.equals(key)) {
					return n;
				}
			}
			return null;
		}
	}

	public V search(K key) {
		Node result = searchForNode(key);
		if (result == null) {
			return null;
		} else {
			return result.value;
		}
	}

	public Collection<V> searchLess(K key) {
		Node foundElem = searchForNode(key);
		if (foundElem == null) {
			foundElem = new Node(null, key);
		}	
		Collection<V> res = new LinkedList<V>();
		for (Node leaf : leafs) {
			if (leaf.key.compareTo(key) < 0) {
				res.add(leaf.value);
			}
		}
		return res;
	}

	public Collection<V> searchLessOrEqual(K key) {
		Node foundElem = searchForNode(key);
		if (foundElem == null) {
			foundElem = new Node(null, key);
		}	
		Collection<V> res = new LinkedList<V>();
		for (Node leaf : leafs) {
			if (leaf.key.compareTo(key) <= 0) {
				res.add(leaf.value);
			}
		}
		return res;
	}

	public Collection<V> searchMoreOrEqual(K key) {
		Node foundElem = searchForNode(key);
		if (foundElem == null) {
			foundElem = new Node(null, key);
		}	
		Collection<V> res = new LinkedList<V>();
		for (Node leaf : leafs) {
			if (leaf.key.compareTo(key) >= 0) {
				res.add(leaf.value);
			}
		}
		return res;
	}

	public Collection<V> searchMore(K key) {
		Node foundElem = searchForNode(key);
		if (foundElem == null) {
			foundElem = new Node(null, key);
		}	
		Collection<V> res = new LinkedList<V>();
		for (Node leaf : leafs) {
			if (leaf.key.compareTo(key) > 0) {
				res.add(leaf.value);
			}
		}
		return res;
	}

	public Collection<V> searchNonEqual(K key) {
		Collection<V> res = new LinkedList<V>();
		Node newNode = new Node(null, key);
		for (Node node : leafs) {
			if (!node.equals(newNode)) {
				res.add(node.value);
			}
		}
		return res;
	}

	private void insertNodeDirect(Node where, Node what) {
		for (int i = 0; i < where.entries.size(); ++i) {
			if (what.key.compareTo(where.entries.get(i).key) == 0) {
				where.entries.remove(i);
				where.entries.add(i, what);
				return;
			}
			if (what.key.compareTo(where.entries.get(i).key) < 0) {
				where.entries.add(i, what);
				where.key = where.entries.get(0).key;
				return;
			}
		}
		where.entries.add(what);
		leafs.add(what);
	}

	public void insert(K key, V data) {
		Node insertNode = new Node(data, key);
		Node parent = searchForParentNode(key);
		if (parent == null) {
			if (root == null) {
				// the tree was empty
				root = insertNode;
				leafs.add(insertNode);
				return;
			} else {
				Node newRoot = new Node();
				newRoot.entries.add(root);
				newRoot.entries.add(insertNode);
				root.parent = newRoot;
				insertNode.parent = newRoot;
				newRoot.key = root.key;
				root = newRoot;
				leafs.add(insertNode);
				return;
			}
		}
		insert(parent, insertNode);
	}

	private void insert(Node where, Node what) {
		if (where.entries.size() < bucketSize) {
			insertNodeDirect(where, what);
		} else {
			Node newNode = new Node();
			for (int i = bucketSize - 1; i > bucketSize / 2; --i) {
				newNode.entries.add(where.entries.remove(i));
			}
			Collections.sort(newNode.entries);
			newNode.key = newNode.entries.get(0).key;
			if (newNode.key.compareTo(what.key) <= 0) {
				insertNodeDirect(newNode, what);
			} else {
				insertNodeDirect(where, what);
			}
			if (where == root) {
				Node newRoot = new Node();
				newRoot.key = root.key;
				newRoot.entries.add(where);
				newRoot.entries.add(newNode);
				newNode.parent = newRoot;
				where.parent = newRoot;
				root = newRoot;
			} else {
				newNode.parent = where.parent;
				insert(where.parent, newNode);
			}
			if (where.isLeaf()) {
				leafs.add(where);
			}
			if (newNode.isLeaf()) {
				leafs.add(newNode);
			}
		}
	}

	private void compactify(Node start) {
		if (start.entries.size() <= bucketSize / 2) {
			// we're underflowed, grab some neighbour
			for (int nNum = 0; nNum < start.parent.entries.size(); ++nNum) {
				Node neighbour = start.parent.entries.get(nNum);
				if (neighbour != start) {
					// either borrow some elements or merge
					if (neighbour.entries.size() <= bucketSize / 2) {
						// merge
						while (start.entries.size() > 0) {
							neighbour.entries.add(start.entries.remove(0));
						}
						Collections.sort(neighbour.entries);
						neighbour.key = neighbour.entries.get(0).key;
						// now start is a valueless leaf, delete it
						start.parent.removeByKey(start.key);
					} else {
						// borrow elements
						int elementsToBorrow = neighbour.entries.size()
								- bucketSize / 2;
						int stopSize = neighbour.entries.size()
								- elementsToBorrow - 1;
						if (start.isLeaf()
								|| neighbour.key.compareTo(start.key) < 0) {
							// take tail of neighbour's elements
							for (int i = neighbour.entries.size() - 1; i > stopSize; --i) {
								start.entries.add(neighbour.entries.remove(i));
							}
						} else {
							// take head of neighbour's elements
							for (int i = 0; i < elementsToBorrow; ++i) {
								start.entries.add(neighbour.entries.remove(0));
							}
							neighbour.key = neighbour.entries.get(0).key;
						}
						Collections.sort(start.entries);
						start.key = start.entries.get(0).key;
					}
					return;
				}
			}
		}
	}

	public K delete(K key) {
		Node parent = searchForParentNode(key);
		if (parent == null) {
			// it's root and root is a leaf
			K res = root.key;
            leafs.remove(root);
            root = null;
			return res;
		} else {
			// it's not root
			Node n = parent.removeByKey(key);
            K res = n.key;
			if (parent != root) {
				compactify(parent);
			}
			leafs.remove(n);
			return res;
		}
	}

	@Override
	public String toString() {
		if (root != null) {
			return root.constructString(0);
		} else {
			return "Empty tree";
		}
	}
}
